{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fi0-b4FfOFDm"
   },
   "source": [
    "##**Pose Estimation with ICP**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VsbCH_XUJDCN"
   },
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the manipulation repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  \n",
    "\n",
    "More details are available [here](http://manipulation.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0pI7YvFePCOG"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
    "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
    "                \"setup_manipulation_colab.py\")\n",
    "    from setup_manipulation_colab import setup_manipulation\n",
    "    setup_manipulation(manipulation_sha='db19656799c11a58bc37dd20604ad38eafb09c62', drake_version='20200921', drake_build='nightly')\n",
    "\n",
    "# Determine if this notebook is currently running as a notebook or a unit test.\n",
    "from IPython import get_ipython\n",
    "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
    "\n",
    "# Setup rendering (with xvfb), if necessary:\n",
    "import os\n",
    "if 'google.colab' in sys.modules and os.getenv(\"DISPLAY\") is None:\n",
    "    from pyvirtualdisplay import Display\n",
    "    display = Display(visible=0, size=(1400, 900))\n",
    "    display.start()\n",
    "\n",
    "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "server_args = []\n",
    "if 'google.colab' in sys.modules:\n",
    "    server_args = ['--ngrok_http_tunnel']\n",
    "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
    "\n",
    "import numpy as np\n",
    "import meshcat\n",
    "import open3d as o3d\n",
    "import meshcat.geometry as g\n",
    "import meshcat.transformations as tf\n",
    "\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph, BaseField, Box, ConnectMeshcatVisualizer, DepthCameraProperties, \n",
    "    DepthImageToPointCloud, DiagramBuilder, FindResourceOrThrow, MakeRenderEngineVtk, CameraInfo, PixelType, \n",
    "    MeshcatPointCloudVisualizer, Parser, RenderEngineVtkParams, RgbdSensor, RigidTransform, Simulator,\n",
    "    RollPitchYaw, RotationMatrix, SpatialInertia, LeafSystem, AbstractValue, PointCloud, Fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mcSz0hOVtBd5"
   },
   "source": [
    "## Problem Description\n",
    "Last lecture, we designed pick and place trajectories **assuming** that the object pose ${}^W X^O$ was known. With all the tools we have learned for goemetric perception, it is time to relax this assumption and finally do pose estimation from sensor data. \n",
    "\n",
    "The goal of the exercise is to give you some real-world experience into what dealing with depth cameras, and what it takes to go from a real depth image to the clean ICP formulation we learned.\n",
    "\n",
    "**These are the main steps of the exercise:**\n",
    "1. Perform Segmentation on the raw pointcloud of the scene to extract pointcloud from the object.\n",
    "2. Tune an off-the-shelf ICP solver and estimate the pose of the object."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NowURj8YP9qb"
   },
   "source": [
    "Before jumping into the main exercise, how should we computationally represent a pointcloud? If we say that pointcloud has $N$ points, then each point has a position in 3D, ${}^Cp^i$, as well as an associated color. Throughout this exercise, we will tend to store them as separate arrays of:\n",
    "- `Nx3` numpy array where each row stores the XYZ position of the point in meters.\n",
    "- `Nx3` numpy array where each row stores the RGB information of the point in `uint8` format. \n",
    "\n",
    "But many opensource libraries have their own implementations of the `PointCloud` class, mainly because they have highly optimized the underlying operations that take place under the hood, usually at the `C/C++`-level. [Drake](https://drake.mit.edu/) has a [PointCloud](https://drake.mit.edu/pydrake/pydrake.perception.html#pydrake.perception.PointCloud) class, and [Open3d](http://www.open3d.org/) has its own [PointCloud](http://www.open3d.org/docs/release/python_api/open3d.geometry.PointCloud.html) class. We will need to switch between different representations, so below here are some helper functions to get us started. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RIFvv4bjMoJd"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Point Cloud Utilities. Converts between three forms of pointcloud objects:\n",
    "  - Drake Pointcloud Object\n",
    "  - Open3d Pointcloud Object\n",
    "  - Nx3 numpy array\n",
    "\"\"\"\n",
    "\n",
    "def pcl_np2o3d(pcl_np):\n",
    "  \"\"\"\n",
    "  Input: Nx3 np array of points.\n",
    "  Ouput: Open3d np object. \n",
    "  \"\"\"\n",
    "  assert(pcl_np.shape[1] == 3) # sanity check\n",
    "  pcl_o3d = o3d.geometry.PointCloud()\n",
    "  pcl_o3d.points = o3d.utility.Vector3dVector(pcl_np)\n",
    "  return pcl_o3d \n",
    "\n",
    "def pcl_o3d2np(pcl_o3d):\n",
    "  \"\"\"\n",
    "  Input: Open3d np object.\n",
    "  Output: Nx3 np array of points. \n",
    "  \"\"\"\n",
    "  return np.asarray(pcl_o3d.points)\n",
    "  \n",
    "def pcl_np2drake(pcl_np, color):\n",
    "  \"\"\"\n",
    "  Input: Nx3 np array of points in mm, and and 3x1 np array of color using uin8 format.\n",
    "  Output: drake Pointcloud object. \n",
    "  \"\"\"\n",
    "  assert(pcl_np.shape[1] == 3) # sanity check. \n",
    "  pcl_drake = PointCloud(new_size = pcl_np.shape[0],\n",
    "                         fields= Fields(BaseField.kXYZs | BaseField.kRGBs))\n",
    "  xyzs = pcl_drake.mutable_xyzs()\n",
    "  xyzs[:,:] = np.array(pcl_np).transpose()\n",
    "  rgbs = pcl_drake.mutable_rgbs()\n",
    "  rgbs[:,:] = np.tile(color, (pcl_np.shape[0], 1)).transpose()\n",
    "  return pcl_drake\n",
    "\n",
    "def pcl_drake2np(pcl_drake):\n",
    "  \"\"\"\n",
    "  Input: drake Pointcloud object.\n",
    "  Output: Nx3 np array of points. \n",
    "  \"\"\"\n",
    "  xyzs = pcl_drake.mutable_xyzs().copy().transpose()\n",
    "  return xyzs \n",
    "\n",
    "def rgb_drake2np(pcl_drake):\n",
    "  rgbs = pcl_drake.mutable_rgbs().copy().transpose()\n",
    "  return rgbs \n",
    "\n",
    "def visualize_pointcloud(vis, pcl, name=\"point_cloud\", X_WC=RigidTransform()):\n",
    "  \"\"\"\n",
    "  pcl: Drake pointcloud object. \n",
    "  All other formats should convert to this format before passing to visualization. \n",
    "  \"\"\"\n",
    "  vis[name].set_object(g.PointCloud(X_WC.multiply(pcl.xyzs()), pcl.rgbs() / 255., size=0.001))\n",
    "\n",
    "def visualize_pointcloud_xyzrgb(vis, pcl_xyz, pcl_rgb, name=\"point_cloud\", X_WC=RigidTransform()):\n",
    "  \"\"\"\n",
    "  Visualization function to access Meshcat's method more directly. \n",
    "  \"\"\"\n",
    "  vis[name].set_object(g.PointCloud(X_WC.multiply(pcl_xyz), pcl_rgb / 255., size=0.001))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e9fQyeITSQkD"
   },
   "source": [
    "## Getting a Pointcloud of the Model ##\n",
    "\n",
    "Before taking a pointcloud of the **scene**, we will need a pointcloud of the **model** to compare against. Generally, this can be done by using existing tools that convert 3D representations (meshes, signed distance functions, etc.) into pointclouds.  \n",
    "\n",
    "Since our red foam brick is of rectangular shape, we'll cheat a bit and generate the points procedurally. When you click the cell below, you should be able to see the red brick and our pointcloud representation of the brick as blue dots. \n",
    "\n",
    "We will save the model pointcloud in the variable `model_pcl_np`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kfyNbppxqGhr"
   },
   "outputs": [],
   "source": [
    "def visualize_red_foam_brick():\n",
    "  \"\"\"\n",
    "  Visualize red foam brick in Meshcat.\n",
    "  \"\"\"\n",
    "  builder = DiagramBuilder()\n",
    "  plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "  parser = Parser(plant)\n",
    "  parser.AddModelFromFile(FindResourceOrThrow(\n",
    "        \"drake/examples/manipulation_station/models/061_foam_brick.sdf\"))\n",
    "  plant.Finalize()\n",
    "\n",
    "  # Setup Meshcat\n",
    "  frames_to_draw = {\"foam_brick\": \"base_link\"}\n",
    "  meshcat = ConnectMeshcatVisualizer(builder, scene_graph, \n",
    "                                     zmq_url=zmq_url, \n",
    "                                     delete_prefix_on_load=False,\n",
    "                                     frames_to_draw=frames_to_draw)\n",
    "  meshcat.load()\n",
    "  builder.Build()\n",
    "  return meshcat.vis\n",
    "\n",
    "def generate_model_pointcloud(xrange, yrange, zrange, res):\n",
    "  \"\"\"\n",
    "  Procedurally generate pointcloud of a rectangle for each side. \n",
    "  \"\"\"\n",
    "  # Decide on how many samples \n",
    "  x_lst = np.linspace(xrange[0], xrange[1], int((xrange[1] - xrange[0]) / res))\n",
    "  y_lst = np.linspace(yrange[0], yrange[1], int((yrange[1] - yrange[0]) / res))\n",
    "  z_lst = np.linspace(zrange[0], zrange[1], int((zrange[1] - zrange[0]) / res))\n",
    "\n",
    "  pcl_lst = [] \n",
    "  # Do XY Plane \n",
    "  for x in x_lst:\n",
    "    for y in y_lst:\n",
    "      pcl_lst.append([x, y, zrange[0]])\n",
    "      pcl_lst.append([x, y, zrange[1]])\n",
    "\n",
    "  # Do YZ Plane\n",
    "  for y in y_lst:\n",
    "    for z in z_lst:\n",
    "      pcl_lst.append([xrange[0], y, z])\n",
    "      pcl_lst.append([xrange[1], y, z])\n",
    "\n",
    "  # Do XZ Plane\n",
    "  for x in x_lst:\n",
    "    for z in z_lst:\n",
    "      pcl_lst.append([x, yrange[0], z])\n",
    "      pcl_lst.append([x, yrange[1], z])\n",
    "\n",
    "\n",
    "  return np.array(pcl_lst)\n",
    "\n",
    "vis = visualize_red_foam_brick()\n",
    "model_pcl_np = generate_model_pointcloud([-0.0375, 0.0375], [-0.025, 0.025], [0., 0.05], 0.002)\n",
    "visualize_pointcloud(vis, pcl_np2drake(model_pcl_np, [0, 0, 255]), \"pcl_scene\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "g6yaPL78TUhD"
   },
   "source": [
    "## Getting the Scene Pointcloud## \n",
    "\n",
    "Now let's set up the ClutteringStation from last lecture and actually take a pointcloud snapshot of the scene with the `red_foam_brick`. We'll place the camera where we have good coverage of the bin. We'll also take a pointcloud snapshot without the `red_foam_brick` so that we can use it for segmentation later.\n",
    "\n",
    "NOTE: There are around `3e7` points that are trying to be published to the visualizer, so things might load slowly, and occasionally the Colab session might crash. Keep calm and run the cells from the beginning! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lA_9j_2pSsTY"
   },
   "outputs": [],
   "source": [
    "vis.delete()\n",
    "\n",
    "def setup_clutter_station(X_WO, X_WC):\n",
    "  builder = DiagramBuilder()\n",
    "  \n",
    "  plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "  parser = Parser(plant)\n",
    "\n",
    "  # Add the foam brick. \n",
    "  if (X_WO is not None):\n",
    "    brick = parser.AddModelFromFile(FindResourceOrThrow(\n",
    "        \"drake/examples/manipulation_station/models/061_foam_brick.sdf\"))\n",
    "    plant.WeldFrames(plant.world_frame(),\n",
    "                    plant.GetFrameByName(\"base_link\", brick),\n",
    "                    X_WO)\n",
    "  \n",
    "  bin1 = parser.AddModelFromFile(FindResourceOrThrow(\n",
    "      \"drake/examples/manipulation_station/models/bin.sdf\"), \"bin1\")\n",
    "  plant.WeldFrames(plant.world_frame(),\n",
    "                   plant.GetFrameByName(\"bin_base\", bin1),\n",
    "                   RigidTransform(RollPitchYaw(0, 0, np.pi/2), [-0.145, -0.63, 0.075]))\n",
    "\n",
    "  bin2 = parser.AddModelFromFile(FindResourceOrThrow(\n",
    "      \"drake/examples/manipulation_station/models/bin.sdf\"), \"bin2\")\n",
    "  plant.WeldFrames(plant.world_frame(),\n",
    "                   plant.GetFrameByName(\"bin_base\", bin2),\n",
    "                   RigidTransform(RollPitchYaw(0, 0, np.pi), [0.5, -0.1, 0.075]))\n",
    "\n",
    "  # Add a rendering engine\n",
    "  renderer = \"my_renderer\"\n",
    "  scene_graph.AddRenderer(renderer,\n",
    "                          MakeRenderEngineVtk(RenderEngineVtkParams()))\n",
    "  \n",
    "  # Add a mock camera model  \n",
    "  camera_instance = plant.AddModelInstance(\"cameras\")\n",
    "  camera = plant.AddRigidBody(\"camera\", camera_instance, SpatialInertia())\n",
    "  plant.RegisterVisualGeometry(\n",
    "      camera, RigidTransform([0, 0, -0.01]), Box(width=.1, depth=.02, height=.02), \"D415\", \n",
    "      [.4, .4, .4, 1.])\n",
    "  plant.WeldFrames(plant.world_frame(), camera.body_frame(), X_WC) \n",
    "  plant.Finalize() \n",
    "\n",
    "  frames_to_draw = {\"cameras\": \"camera\"}\n",
    "  meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url, \n",
    "                                      delete_prefix_on_load=False,\n",
    "                                      frames_to_draw=frames_to_draw)\n",
    "  meshcat.load() \n",
    "\n",
    "  properties = DepthCameraProperties(width=640,\n",
    "                                      height=480,\n",
    "                                      fov_y=np.pi / 4.0,\n",
    "                                      renderer_name=renderer,\n",
    "                                      z_near=0.1,\n",
    "                                      z_far=10.0)\n",
    "  camera = builder.AddSystem(\n",
    "      RgbdSensor(parent_id=plant.GetBodyFrameIdOrThrow(camera.index()),\n",
    "                  X_PB=RigidTransform(),\n",
    "                  properties=properties,\n",
    "                  show_window=False))\n",
    "  camera.set_name(\"rgbd_sensor\")\n",
    "  builder.Connect(scene_graph.get_query_output_port(),\n",
    "                  camera.query_object_input_port())\n",
    "\n",
    "  # Export the camera outputs\n",
    "  builder.ExportOutput(camera.color_image_output_port(), \"color_image\")\n",
    "  builder.ExportOutput(camera.depth_image_32F_output_port(), \"depth_image\")\n",
    "\n",
    "  # Add a system to convert the camera output into a point cloud\n",
    "  to_point_cloud = builder.AddSystem(\n",
    "      DepthImageToPointCloud(camera_info=camera.depth_camera_info(),\n",
    "                              pixel_type=PixelType.kDepth32F,\n",
    "                              fields=BaseField.kXYZs | BaseField.kRGBs))\n",
    "  builder.Connect(camera.depth_image_32F_output_port(),\n",
    "                  to_point_cloud.depth_image_input_port())\n",
    "  builder.Connect(camera.color_image_output_port(),\n",
    "                  to_point_cloud.color_image_input_port())\n",
    "\n",
    "  # Send the point cloud to meshcat for visualization, too.\n",
    "  meshcat_pointcloud = builder.AddSystem(MeshcatPointCloudVisualizer(meshcat, X_WP=X_WC, draw_period=1./5.))\n",
    "  builder.Connect(to_point_cloud.point_cloud_output_port(), meshcat_pointcloud.get_input_port())\n",
    "\n",
    "  # Export the point cloud output.\n",
    "  builder.ExportOutput(to_point_cloud.point_cloud_output_port(),\n",
    "                        \"point_cloud\")\n",
    "\n",
    "  diagram = builder.Build()\n",
    "  diagram.set_name(\"depth_camera_demo_system\")\n",
    "  return diagram \n",
    "\n",
    "# Set pose of the brick \n",
    "X_WO = RigidTransform(RollPitchYaw(0, 0, np.pi/5).ToRotationMatrix(), np.array([-0.1, -0.6, 0.09]))\n",
    "\n",
    "# Setup CameraPose\n",
    "X_WC = RigidTransform(\n",
    "    RollPitchYaw(0, 0, 0).ToRotationMatrix().multiply(\n",
    "        RollPitchYaw(-np.pi/2. - np.pi/3, 0, 0).ToRotationMatrix()),\n",
    "    [-.1, -.8, .5])\n",
    "\n",
    "# Take a pointcloud snapshot of the background to use for subtraction \n",
    "diagram = setup_clutter_station(None, X_WC)\n",
    "simulator = Simulator(diagram)\n",
    "simulator.AdvanceTo(0.01)\n",
    "context = simulator.get_context()\n",
    "scene_pcl_drake_background =  diagram.GetOutputPort(\"point_cloud\").Eval(context)\n",
    "scene_pcl_np_background = X_WC.multiply(pcl_drake2np(scene_pcl_drake_background).transpose()).transpose()\n",
    "scene_rgb_np_background = rgb_drake2np(scene_pcl_drake_background)\n",
    "\n",
    "# Take a pointcloud snapshot of the scene with the brick. \n",
    "diagram = setup_clutter_station(X_WO, X_WC)\n",
    "simulator = Simulator(diagram)\n",
    "simulator.AdvanceTo(0.01)\n",
    "context = simulator.get_context()\n",
    "scene_pcl_drake = diagram.GetOutputPort(\"point_cloud\").Eval(context)\n",
    "scene_pcl_np = X_WC.multiply(pcl_drake2np(scene_pcl_drake).transpose()).transpose()\n",
    "scene_rgb_np = rgb_drake2np(scene_pcl_drake)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MbgWqzgcUaAX"
   },
   "source": [
    "## Visualizing the Problem ##\n",
    "\n",
    "That was a lot of work, but if you run the below cell, Meshcat will finally show you a clean formulation of the main problem. We have 3 pointcloud objects in Meshcat:\n",
    "\n",
    "- `pcl_model`: Pointcloud of models\n",
    "- `pcl_scene`: Raw pointcloud of the foam-brick scene obtained from a RGBD camera.\n",
    "- `pcl_scene_background`: Raw pointcloud of the background obtained from a RGBD camera. \n",
    "\n",
    "In case you forgot, In Meshcat's menu you can go into the `meshcat` tab and turn different objects on and off so that you can see what the background pointcloud looks like as well. \n",
    "\n",
    "NOTE: You might have to wait a bit until the bin pointcloud shows up.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0zMWmIMh5upv"
   },
   "outputs": [],
   "source": [
    "vis = meshcat.Visualizer(zmq_url)\n",
    "vis.delete()\n",
    "\n",
    "visualize_pointcloud(vis, pcl_np2drake(model_pcl_np, [0, 0, 255]), name=\"pcl_model\")\n",
    "visualize_pointcloud(vis, scene_pcl_drake, name=\"pcl_scene\", X_WC=X_WC)\n",
    "visualize_pointcloud_xyzrgb(vis, \n",
    "                            scene_pcl_np_background.transpose(), \n",
    "                            scene_rgb_np_background.transpose(), \n",
    "                            name=\"pcl_scene_background\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TriejkVOWlsN"
   },
   "source": [
    "If we simply run ICP with `pcl_model` and `pcl_scene`, we might get a terrible result because there might be features in the background that the model is trying to run correspondence with. So we'd like to vet the problem a bit and perform **segmentation**: which parts of the scene pointcloud corresponds to an actual point on the `red_foam_brick`? \n",
    "\n",
    "\n",
    "**Now it's your turn to code!**\n",
    "\n",
    "Below, you will implement a function `segment_scene_pcl` that takes in a pointcloud of the scene and return the relevant points that are actually on the `red_foam_brick`. But here are the rules of the game:\n",
    "- You **may** use color data, the background pointcloud, and any outlier detection algorithm that you can write to perform segmentation.\n",
    "- You may **not** explicitly impose conditions on the position to filter out the data. Remember that our goal is to estimate the pose in the first place, so using position will be considered cheating.\n",
    "- You may **not** use external libraries that are not in this notebook already. \n",
    "\n",
    "In order to get full score for this assignment, you need to satisfy both criteria:\n",
    "- The number of false outliers (points which are not on the red brick but was caught by segmentation) must not exceed 80 points.\n",
    "- The number of missed inliers (points that are on the red brick but was not caught by segmentation) must not exceed 80 points. \n",
    "\n",
    "You will be able to visualize your segmented pointclouds on Meshcat by running the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2OY3fQjJGU3A"
   },
   "outputs": [],
   "source": [
    "def segment_scene_pcl(scene_pcl_np, scene_rgb_np, scene_pcl_np_background, scene_rgb_np_background):\n",
    "  \"\"\"\n",
    "  Inputs: \n",
    "  scene_pcl_np: Nx3 np.float32 array of pointclouds, each row containing xyz position of each point in meters. \n",
    "  scene_rgb_np: Nx3 np.uint8   array of pointclouds, each row containing rgb color data of each point.\n",
    "  scene_pcl_np_background: Nx3 np.float32 array of pointclouds, each row containing xyz position of each point in meters. \n",
    "  scene_rgb_np_background: Nx3 np.uint8   array of pointclouds, each row containing rgb color data of each point.\n",
    "\n",
    "  Outputs:\n",
    "  scene_pcl_np_filtered: Mx3 np.float32 array of pointclouds that are on the foam brick. \n",
    "  \"\"\"\n",
    "  ####################\n",
    "  # Fill your code here. \n",
    "\n",
    "  scene_pcl_np_filtered = scene_pcl_np\n",
    "  ####################\n",
    " \n",
    "  return scene_pcl_np_filtered\n",
    "\n",
    "scene_pcl_np_filtered = segment_scene_pcl(scene_pcl_np,\n",
    "                                          scene_rgb_np,\n",
    "                                          scene_pcl_np_background,\n",
    "                                          scene_rgb_np_background)\n",
    "visualize_pointcloud(vis, pcl_np2drake(scene_pcl_np_filtered, [0, 255, 0]), \n",
    "                     name='pcl_scene_filtered')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FiZex_EDa-RC"
   },
   "source": [
    "## ICP for Pose Estimation\n",
    "\n",
    "Now that we have a subset of scene points that we want to use to estimate the pose, let's do ICP to figure out what ${}^W X^O$ is. Instead of implementing your own ICP this time, we will use the most common implementation available: it comes from [Open3d](http://www.open3d.org/) and implements the method from [[Rusinkiewicz, 2001]](http://www.pcl-users.org/file/n4037867/Rusinkiewicz_Effcient_Variants_of_ICP.pdf).\n",
    "\n",
    "We know that ICP can't work very well without even a rough initialization. Let's assume that we at least know that the `red_foam_brick` is inside the bin, so that we can initialize the ${}^W X^O$ to be at the center of the bin with an identity rotation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XADq9uCGK2rV"
   },
   "outputs": [],
   "source": [
    "initial_guess = RigidTransform()\n",
    "initial_guess.set_translation([-0.145, -0.63, 0.09]) # do not modify \n",
    "\n",
    "def o3d_icp(model_pcl_np, scene_pcl_np_filtered, initial_guess, solver_params): \n",
    "  \"\"\"Performs ICP with Open3d's Generalized-ICP implementation.\n",
    "  Input: \n",
    "    - model_pcl_np: Nx3 np.float32 array, model pointcloud \n",
    "    - scene_pcl_np_filtered: Nx3 np.float32 array, model pointcloud\n",
    "    - initial_guess: Drake RigidTransform object used as initial guess to correspondence. \n",
    "    - solver_params: dictionary containing solver parameters.\n",
    "  Output:\n",
    "    - X_MS: Estimated relative transform between model and scene pointclouds\n",
    "  \"\"\"\n",
    "  source = pcl_np2o3d(model_pcl_np)\n",
    "  target = pcl_np2o3d(scene_pcl_np_filtered)\n",
    "\n",
    "  reg_p2p = o3d.registration.registration_icp(\n",
    "      source, target, \n",
    "      solver_params[\"max_correspondence_distance\"], initial_guess.GetAsMatrix4(),\n",
    "      estimation_method = o3d.registration.TransformationEstimationPointToPoint(),\n",
    "      criteria = o3d.registration.ICPConvergenceCriteria(\n",
    "          relative_fitness=solver_params[\"relative_fitness\"],\n",
    "          relative_rmse=solver_params[\"relative_rmse\"],\n",
    "          max_iteration=solver_params[\"max_iteration\"])\n",
    "  )\n",
    "\n",
    "  X_MS_hat = RigidTransform(reg_p2p.transformation)\n",
    "\n",
    "  return X_MS_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Cr4AkNyohLYF"
   },
   "source": [
    "You will notice that Open3d's ICP method exposes some parameters you need to tune for it to be working successfully. The [documentation](http://www.open3d.org/docs/release/python_api/open3d.registration.registration_icp.html?highlight=registration_icp) describes the parameters as follows: \n",
    "\n",
    "- `max_correspondence_distance`: maximum distance for pair of points to be registered as a correspondence. \n",
    "- `relative_fitness`: If relative change (difference) of fitness score is lower than `relative_fitness`, the iteration stops. \n",
    "- `relative_rmse`: If relative change (difference) of inlier RMSE score is lower than `relative_rmse`, the iteration stops. \n",
    "- `max_iteration`: maximum iteration before iteration stops.\n",
    "\n",
    "**Now, it's your time to code**\n",
    "\n",
    "The problem is quite straightforward - if you had a good segmented pointcloud for the previous question, can you tweak the solver to find a good estimate of ${^WX^O}$? \n",
    "\n",
    "You should be able to check the validity of your solution on Meshcat by running the cell below.\n",
    "\n",
    "Your estimated transform is stored in a variable `X_MS_hat`, which we will be used for the grading code. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "n-wijftedyAH"
   },
   "outputs": [],
   "source": [
    "# Modify this dictionary. \n",
    "solver_params = {\n",
    "    \"max_correspondence_distance\": 1.0,\n",
    "    \"relative_fitness\": 1e-6,\n",
    "    \"relative_rmse\": 1e-6,\n",
    "    \"max_iteration\": 30\n",
    "}\n",
    "\n",
    "X_MS_hat = o3d_icp(model_pcl_np, scene_pcl_np_filtered, initial_guess, solver_params)\n",
    "visualize_pointcloud(vis, pcl_np2drake(model_pcl_np, [255, 0, 255]), name=\"pcl_estimated\", X_WC=X_MS_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "O7XDta3wyO2Z"
   },
   "source": [
    "In order to get full points, the relative transform ${}^OX^{\\hat{O}}$, between the true pose ${}^W X^O$ and estimated pose ${}^WX^\\hat{O}$, must satisfy the following criteria: \n",
    "- xyz translation values of ${}^O X^{\\hat{O}}$ does not exceed 0.005 (5 milimeters)\n",
    "- RollPitchYaw angles of ${}^O X^{\\hat{O}}$ does not exceed 0.05 (about 3degs)\n",
    "\n",
    "We will give partial scores to translations of less than 0.02 (2 centimeters), and rotations less than 0.05 (~30 degrees)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rcZCfttlbufD"
   },
   "outputs": [],
   "source": [
    "np.set_printoptions(precision=3, suppress=True)\n",
    "X_OOhat = X_MS_hat.inverse().multiply(X_WO)\n",
    "\n",
    "rpy = RollPitchYaw(X_OOhat.rotation()).vector()\n",
    "xyz = X_OOhat.translation()\n",
    "\n",
    "print(\"RPY Value: \" + str(rpy))\n",
    "print(\"XYZ Value: \" + str(xyz))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zPmeRLtJk410"
   },
   "source": [
    "##How will this notebook be Graded?##\n",
    "\n",
    "If you are enrolled in the class, this notebook will be graded using [Gradescope](www.gradescope.com). You should have gotten the enrollement code on our announcement in Piazza. \n",
    "\n",
    "For submission of this assignment, you must do as follows:. \n",
    "- Download and submit the notebook `pick_and_place_perception.ipynb` to Gradescope's notebook submission section, along with your notebook for the other problems.\n",
    "\n",
    "We will evaluate the local functions in the notebook to see if the function behaves as we have expected. For this exercise, the rubric is as follows:\n",
    "- [4 pts] `segment_scene_pcl` correctly segments the scene by having less than 80 missed inliers and 80 false outliers.\n",
    "- [1 pts] XYZ Translation of ${}^O X^{\\hat{O}}$ does not exceed 0.02. \n",
    "- [1 pts] XYZ Translation of ${}^O X^{\\hat{O}}$ does not exceed 0.005. \n",
    "- [1 pts] RollPitchYaw angles of ${}^O X^{\\hat{O}}$ does not exceed 0.5. \n",
    "- [1 pts] RollPitchYaw angles of ${}^O X^{\\hat{O}}$ does not exceed 0.05. \n",
    "\n",
    "Unlike other exercises, there are partial points for this exercise depending on how much you've violated these requirements. \n",
    "\n",
    "Below is our autograder where you can check your score!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WrTSDUHk3S_J"
   },
   "outputs": [],
   "source": [
    "from manipulation.exercises.pose.test_pose_estimation import TestPoseEstimation\n",
    "from manipulation.exercises.grader import Grader \n",
    "\n",
    "Grader.grade_output([TestPoseEstimation], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "colab": {
   "collapsed_sections": [],
   "name": "pose_estimation_icp.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}